from ase import Atoms
from quippy.descriptors import Descriptor
from quippy.potential import Potential
#from ase.build import bulk
#from ase.build import nanotube

from core_atomistic.atomic_model import AtomicModel


class QuipCalc(object):
    """The atom class."""

    def __init__(self):
        """Constructor"""
        self.pot: Potential = None
        self.desc: Descriptor = None

    def set_potential(self, calc_args='Potential xml_label="GAP_2024_5_28_300_2_25_15_266"',
                      filename='/home/sozykin/calc1tb/GAPs/probe07/GAP.xml'):
        print("init_args ", calc_args)
        print("param_filename  ", filename)
        self.pot = Potential(init_args=calc_args,
                             param_filename=filename)

    def potential(self, model: AtomicModel):
        # construct a potential, usage is the same as `quip` command line
        #pot = Potential('xml_label=GAP_2024_3_19_300_11_26_3_500',
        #                param_filename='/home/sozykin/calc1tb/gap/GAP.xml',
        #                calc_args='local_gap_variance')

        """
        I'm trying to receive the energies of the models in the quippy. However, the code

```
from quippy.potential import Potential
from ase.build import nanotube
cnt2 = nanotube(3, 3, length=6, bond=1.4, symbol='Si')
Si_potential = Potential(init_args='Potential xml_label="GAP_2017_6_17_60_4_3_56_165"',
                                      param_filename='/home/sozykin/Downloads/Si_PRX_GAP/gp_iter6_sparse9k.xml')
cnt2.set_calculator(Si_potential)
print("Si: ", cnt2.get_potential_energy())
```

only works for some potentials ([https://doi.org/10.17863/CAM.65004](https://doi.org/10.17863/CAM.65004)).
For the rest ([https://github.com/AI4Materials/2DM-GAP/tree/submission](https://github.com/AI4Materials/2DM-GAP/tree/submission), GAP-20), I get nan instead of energy.
What do you think this might be connected with? In the Python code I only change init_args and param_filename.
        """

        #cnt1 = nanotube(6, 0, length=4)
        #cnt2 = nanotube(3, 3, length=6, bond=1.4, symbol='Si')

        #Si_potential = Potential(init_args='Potential xml_label="GAP_2017_6_17_60_4_3_56_165"',
        #                         param_filename='/home/sozykin/Downloads/Si_PRX_GAP/gp_iter6_sparse9k.xml')

        #C_potential = Potential(init_args='Potential xml_label="GAP_2021_4_21_0_21_17_24_635"',
        #                        param_filename='/home/sozykin/Downloads/GAP_20_U_s0.1/Carbon_GAP_V10_2_LDIter1_kptconsistent_grr6_2b_3b_soap_n12l4_sp9k_delta-2-0.05-0.1.xml')

        #C_potential1 = Potential(init_args='Potential xml_label="GAP_2022_2_28_180_5_42_29_949"',
        #                     param_filename='/home/sozykin/Downloads/2DM-GAP-submission/AI4Materials-2DM-GAP-2788796/potentials/Graphene/Gr.xml')

        # self.pot = Potential(init_args='Potential xml_label="GAP_2022_4_3_180_0_7_12_229"',
        #                      param_filename='/home/sozykin/Downloads/2DM-GAP-submission/AI4Materials-2DM-GAP-2788796/potentials/Silicene/Si.xml')

        if self.pot is None:
            print("No potential")
            return

        model = Atoms(model.formula(),
                      positions=model.get_positions(),
                      cell=model.get_cell(),
                      pbc=[1, 1, 1])

        model.set_calculator(self.pot)
        e0 = model.get_potential_energy()
        print(e0)
        return e0

    def set_descriptor(self, l_max: int = 6, n_max=12, cutoff=5.0, atom_sigma=0.5):
        params = ("l_max=" + str(int(l_max)) + " n_max=" + str(int(n_max)) + " cutoff=" + str(cutoff) +
                  " atom_sigma=" + str(atom_sigma))
        self.desc = Descriptor("soap " + params)
        
    def descriptor(self, model: AtomicModel):
        if self.desc is None:
            print("No descriptor")
            return

        model = Atoms(model.formula(),
                      positions=model.get_positions(),
                      cell=model.get_cell(),
                      pbc=[1, 1, 1])
        d = self.desc.calc(model)['data']
        return d

